package kodkod.analysis;

import java.util.ArrayList;
import java.util.HashMap;

import kodkod.analysis.ExpressionAnalyzer.Op;
import kodkod.ast.*;
import kodkod.instance.*;
import kodkod.util.ints.IndexedEntry;
import kodkod.util.ints.IntIterator;
import kodkod.util.ints.IntSet;
import kodkod.util.ints.IntTreeSet;
import kodkod.util.ints.Ints;

public class BoundsExtension {
	private static TupleFactory factory = null;
	private static int base = 0;
	
	public IntSet upper;
	public IntSet lower;
	
	public int arity;
	public int[] max;
	public int[] min;
	
	Node expression;
	
	ArrayList<BoundsExtension> children = new ArrayList<BoundsExtension>();
	enum Op {
		NONE, LEAF, DECL, CONST, CLOSURE, REFLEXIVE_CLOSURE, TRANSPOSE, UNION, INTERSECTION, DIFFERENCE, PRODUCT, JOIN, OVERRIDE, COMPREHENSION, IFTHENELSE, PROJECTION, CAST
	}
	Op type = Op.NONE;
	
	public static HashMap<Expression, BoundsExtension> hash = new HashMap<Expression, BoundsExtension>();
	
	public static BoundsExtension get(Expression e) {
		if (factory != Analyzer.translation.bounds.universe().factory()) {
			factory = Analyzer.translation.bounds.universe().factory();
			base = factory.universe().size();
			hash = new HashMap<Expression, BoundsExtension>();
		}
		
		if (!hash.containsKey(e)) {
			hash.put(e, new BoundsExtension(e));
		}
		return hash.get(e);
	}
	
	public static void purgeHash() {
		hash = new HashMap<Expression, BoundsExtension>();
		System.gc();
	}
	
	private BoundsExtension(Expression expression) {
		this.expression = expression;
		arity = expression.arity();
		min = new int[arity];
		max = new int[arity];
		for (int pos = 0; pos < arity; pos++) {
			min[pos] = base;
			max[pos] = 0;
		}
		upper = new IntTreeSet();
		lower = new IntTreeSet();
		
		if (expression instanceof LeafExpression) {
			if (expression instanceof Relation) {
				type = Op.LEAF;
			} else if (expression instanceof Variable) {
				if (Analyzer.translation.skolemRelations.containsKey(expression)) {
					type = Op.LEAF;
					this.expression = Analyzer.translation.skolemRelations.get(expression);
				} else {
					type = Op.DECL;
					children.add(BoundsExtension.get(Analyzer.decls.get((Variable)expression)));
				}
			} else if (expression instanceof ConstantExpression) {
				type = Op.CONST;
			}
		} else if (expression instanceof UnaryExpression) {
			UnaryExpression exp = (UnaryExpression)expression;
			children.add(BoundsExtension.get(exp.expression()));
			switch (exp.op()) {
				case REFLEXIVE_CLOSURE: type = Op.REFLEXIVE_CLOSURE; break;
				case CLOSURE: 	type = Op.CLOSURE; break;
				case TRANSPOSE: type = Op.TRANSPOSE; break;
			}
		} else if (expression instanceof BinaryExpression) {
			BinaryExpression exp = (BinaryExpression)expression;
			children.add(BoundsExtension.get(exp.left()));
			children.add(BoundsExtension.get(exp.right()));
			switch (exp.op()) {
				case DIFFERENCE: 	type = Op.DIFFERENCE; break;
				case INTERSECTION: 	type = Op.INTERSECTION; break;
				case JOIN: 			type = Op.JOIN; break;
				case OVERRIDE: 		type = Op.OVERRIDE; break;
				case PRODUCT: 		type = Op.PRODUCT; break;
				case UNION: 		type = Op.UNION; break;
			}
		} else if (expression instanceof NaryExpression) {
			NaryExpression exp = (NaryExpression)expression;
			for (Expression child : exp) {
				children.add(BoundsExtension.get(child));
			}
			switch (exp.op()) {
				case INTERSECTION: 	type = Op.INTERSECTION; break;
				case OVERRIDE: 		type = Op.OVERRIDE; break;
				case PRODUCT: 		type = Op.PRODUCT; break;
				case UNION: 		type = Op.UNION; break;
			}
		} else if (expression instanceof Comprehension) {
			Comprehension exp = (Comprehension)expression;
			for (Decl decl : exp.decls()) {
				children.add(BoundsExtension.get(decl.expression()));
			}
			type = Op.COMPREHENSION;
		} else if (expression instanceof IfExpression) {
			IfExpression exp = (IfExpression)expression;
			children.add(BoundsExtension.get(exp.thenExpr()));
			children.add(BoundsExtension.get(exp.elseExpr()));
			type = Op.IFTHENELSE;
		} else if (expression instanceof ProjectExpression) {
			ProjectExpression exp = (ProjectExpression)expression;
			children.add(BoundsExtension.get(exp.expression()));
			type = Op.PROJECTION;
		} else if (expression instanceof IntToExprCast) {
			IntToExprCast exp = (IntToExprCast)expression;
			type = Op.CAST;
		}
		initBounds();
		if (type != Op.CONST && type != Op.DECL && type != Op.CLOSURE && type != Op.REFLEXIVE_CLOSURE) updateMinMax();
	}
	
	public int extractAtom(int index, int atomIndex) {
		int trail = index / bpow(atomIndex);
		return trail % base;
	}
	
	public void updateMinMax(int index) {
		for (int pos = 0; pos < arity; pos++) {
			int atom = extractAtom(index, pos);
			if (atom < min[pos]) min[pos] = atom;
			if (atom > max[pos]) max[pos] = atom;
		}
	}
	
	public void updateMinMax() {
		for (int pos = 0; pos < arity; pos++) {
			min[pos] = base;
			max[pos] = 0;
		}
		IntIterator iter = upper.iterator();
		while (iter.hasNext()) {
			updateMinMax(iter.next());
		}
	}

	public int arity() {
		return arity;
	}
	
	public int size() {
		return upper.size() - lower.size();
	}
	
	private int bpow(int exp) {
		return (int)Math.pow(base, exp);
	}
	
	public void initBounds() {
		switch (this.type) {
		case LEAF:
			upper = Analyzer.translation.bounds.upperBound((Relation)expression).indexView();
			lower = Analyzer.translation.bounds.lowerBound((Relation)expression).indexView();
			break;
		case CONST:
			if (expression.equals(ConstantExpression.UNIV)) {
				lower = upper = factory.allOf(1).indexView();
				min[0] = 0; max[0] = base-1;
			} else if (expression.equals(ConstantExpression.IDEN)) {
				for (int i = 0; i < base; i++) {
					upper.add(i*base+i);
				}
				lower = upper;
				min[0] = 0; max[0] = base-1;
				min[1] = 0; max[1] = base-1;
			} else if (expression.equals(ConstantExpression.NONE)) {
				lower = upper = factory.noneOf(1).indexView();
				min[0] = 0; max[0] = 0;
			} else if (expression.equals(ConstantExpression.INTS)) {
				lower = upper;
				for (IndexedEntry<TupleSet> entry : Analyzer.translation.bounds.intBounds()) {
					for (Tuple t : entry.value()) {
						upper.add(t.index());
						updateMinMax(t.index());
					}
				}
			}
			break;
		case DECL:
			upper = children.get(0).upper;
			min = children.get(0).min; 
			max = children.get(0).max;
			break;
		case TRANSPOSE:
			IntIterator iter = children.get(0).upper.iterator();
			while (iter.hasNext()) {
				int index = iter.next();
				upper.add(index / base + (index % base) * base);
			}
			iter = children.get(0).lower.iterator();
			while (iter.hasNext()) {
				int index = iter.next();
				lower.add(index / base + (index % base) * base);
			}
			break;
		case CLOSURE:
			upper.addAll(children.get(0).upper);
			lower.addAll(children.get(0).lower);
			min = children.get(0).min;
			max = children.get(0).max;
		case REFLEXIVE_CLOSURE: //fallthrough intended
			for (int i = 0; i < base; i++) {
				upper.add(i*base+i);
				lower.add(i*base+i);
			}
			break;
		case IFTHENELSE:
		case UNION:
		case OVERRIDE:
			for (BoundsExtension child : children) {
				upper.addAll(child.upper);
				lower.addAll(child.lower);
			}
			break;
		case INTERSECTION:
			upper.addAll(children.get(0).upper);
			lower.addAll(children.get(0).lower);
			for (int i = 1; i < children.size(); i++) {
				upper.retainAll(children.get(0).upper);
				lower.retainAll(children.get(0).lower);
			}
			break;
		case DIFFERENCE:
			upper.addAll(children.get(0).upper);
			lower.addAll(children.get(0).lower);
			upper.removeAll(children.get(1).lower);
			lower.removeAll(children.get(1).lower);
			break;
		case PRODUCT:
		case COMPREHENSION:
			IntTreeSet intermediate = new IntTreeSet();
			intermediate.addAll(children.get(0).upper);
			IntTreeSet nextIntermediate = new IntTreeSet();
			for (int i = 1; i < children.size(); i++) {
				BoundsExtension child = children.get(i);
				nextIntermediate = new IntTreeSet();
				IntIterator iter1 = intermediate.iterator();
				while (iter1.hasNext()) {
					int index1 = iter1.next();
					IntIterator iter2 = child.upper.iterator();
					while (iter2.hasNext()) {
						int index2 = iter2.next();
						nextIntermediate.add(index1 * bpow(child.arity) + index2);
					}
				}
				intermediate = nextIntermediate;
			}
			upper.addAll(nextIntermediate);
			
			intermediate = new IntTreeSet();
			intermediate.addAll(children.get(0).lower);
			nextIntermediate = new IntTreeSet();
			for (int i = 1; i < children.size(); i++) {
				BoundsExtension child = children.get(i);
				nextIntermediate = new IntTreeSet();
				IntIterator iter1 = intermediate.iterator();
				while (iter1.hasNext()) {
					int index1 = iter1.next();
					IntIterator iter2 = child.lower.iterator();
					while (iter2.hasNext()) {
						int index2 = iter2.next();
						nextIntermediate.add(index1 * bpow(child.arity) + index2);
					}
				}
				intermediate = nextIntermediate;
			}
			lower.addAll(nextIntermediate);
			
			break;
		case JOIN:
			IntIterator iter1 = children.get(0).upper.iterator();
			while (iter1.hasNext()) {
				int index1 = iter1.next();
				IntIterator iter2 = children.get(1).upper.iterator();
				while (iter2.hasNext()) {
					int index2 = iter2.next();
					int c1 = index1 % base;
					int c2 = index2 / bpow(children.get(1).arity-1);
					if (c1 == c2) {
						int index = bpow(children.get(1).arity-1) * (index1 / base) + index2 % bpow(children.get(1).arity-1);
						upper.add(index);
					}
				}
			}
			iter1 = children.get(0).lower.iterator();
			while (iter1.hasNext()) {
				int index1 = iter1.next();
				IntIterator iter2 = children.get(1).lower.iterator();
				while (iter2.hasNext()) {
					int index2 = iter2.next();
					int c1 = index1 % base;
					int c2 = index2 / bpow(children.get(1).arity-1);
					if (c1 == c2) {
						int index = bpow(children.get(1).arity-1) * (index1 / base) + index2 % bpow(children.get(1).arity-1);
						lower.add(index);
					}
				}
			}
			break;
		case PROJECTION:
			ProjectExpression exp = (ProjectExpression)expression;
			ArrayList<Integer> cols = new ArrayList<Integer>();
			for (int col = 0; col < arity; col++) {
				IntExpression intExp = exp.column(col);
				if (intExp instanceof IntConstant) {
					int column = ((IntConstant)intExp).value();
					cols.add(column);
				}
			}
			
			iter = children.get(0).upper.iterator();
			while (iter.hasNext()) {
				int index = iter.next();
				ArrayList<Object> list = new ArrayList<Object>();
				for (int col : cols) {
					list.add(factory.tuple(get(exp.expression()).arity, index).atom(col));
				}
				upper.add(factory.tuple(list).index());
			}
			iter = get(exp.expression()).lower.iterator();
			while (iter.hasNext()) {
				int index = iter.next();
				ArrayList<Object> list = new ArrayList<Object>();
				for (int col : cols) {
					list.add(factory.tuple(get(exp.expression()).arity, index).atom(col));
				}
				upper.add(factory.tuple(list).index());
			}
			break;
		case CAST:
			for (IndexedEntry<TupleSet> entry : Analyzer.translation.bounds.intBounds()) {
				TupleSet ts = entry.value();
				this.upper.addAll(ts.indexView());
				this.lower.addAll(ts.indexView());
			}
			break;
		case NONE:
		default:
			System.err.println("Warning: Unhandled Expression Type @"+this.expression);
		}
	}
}
